#pragma once

#include <stdint.h>
#include "vector/tensor.h"
#include "vector/layer.h"

class Linear : public Layer
{
public:
    Linear(int64_t d, int64_t h);
    virtual ~Linear() = default;

    virtual void forward(const std::vector<std::vector<double>> &input, std::vector<std::vector<double>> &output) override;
    virtual std::vector<std::vector<double>> backward(const std::vector<std::vector<double>> &output_gradient) override;

    std::vector<std::vector<double>> get_weights();
    void set_weights(const std::vector<std::vector<double>> &weights);

    std::vector<double> get_bias();
    void set_bias(const std::vector<double> &bias);

    std::vector<std::vector<double>> get_weights_gradient();
    std::vector<double> get_bias_gradient();
private:
    std::vector<std::vector<double>> input_{};
    std::vector<std::vector<double>> weights_{};
    std::vector<std::vector<double>> weights_gradient_{};
    std::vector<double> bias_{};
    std::vector<double> bias_gradient_{};
};
